from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# 虚拟的电影数据
movies = [
    {"title": "The Shawshank Redemption", "genre": "Drama"},
    {"title": "The Godfather", "genre": "Crime"},
    {"title": "The Dark Knight", "genre": "Action"},
    {"title": "Pulp Fiction", "genre": "Crime"},
    {"title": "Forrest Gump", "genre": "Drama"},
    {"title": "Fight Club", "genre": "Drama"},
    {"title": "Inception", "genre": "Action"},
    {"title": "The Matrix", "genre": "Action"},
    {"title": "Schindler's List", "genre": "Biography"},
    {"title": "The Lord of the Rings: The Return of the King", "genre": "Adventure"},
    {"title": "Forrest Gump", "genre": "Drama"},
    {"title": "The Green Mile", "genre": "Drama"},
    {"title": "The Godfather Part II", "genre": "Crime"},
    {"title": "The Silence of the Lambs", "genre": "Thriller"},
    {"title": "Goodfellas", "genre": "Crime"},
    {"title": "The Shawshank Redemption", "genre": "Drama"},
    {"title": "Saving Private Ryan", "genre": "War"},
    {"title": "The Departed", "genre": "Crime"},
    {"title": "Gladiator", "genre": "Action"},
    {"title": "Titanic", "genre": "Romance"},
]

# 为每部电影添加简短的描述
movie_descriptions = {
    "The Shawshank Redemption": "Two imprisoned men bond over a number of years, finding solace and eventual redemption through acts of common decency.",
    "The Godfather": "The aging patriarch of an organized crime dynasty transfers control of his clandestine empire to his reluctant son.",
    "The Dark Knight": "When the menace known as the Joker wreaks havoc and chaos on the people of Gotham, Batman must accept one of the greatest psychological and physical tests of his ability to fight injustice.",
    "Pulp Fiction": "The lives of two mob hitmen, a boxer, a gangster and his wife, and a pair of diner bandits intertwine in four tales of violence and redemption.",
    "Forrest Gump": "The presidencies of Kennedy and Johnson, the Vietnam War, the Watergate scandal and other historical events unfold from the perspective of an Alabama man with an IQ of 75, whose only desire is to be reunited with his childhood sweetheart.",
    "Fight Club": "An insomniac office worker and a devil-may-care soapmaker form an underground fight club that evolves into something much, much more.",
    "Inception": "A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O.",
    "The Matrix": "A computer hacker learns from mysterious rebels about the true nature of his reality and his role in the war against its controllers.",
    "Schindler's List": "In German-occupied Poland during World War II, industrialist Oskar Schindler gradually becomes concerned for his Jewish workforce after witnessing their persecution by the Nazis.",
    "The Lord of the Rings: The Return of the King": "Gandalf and Aragorn lead the World of Men against Sauron's army to draw his gaze from Frodo and Sam as they approach Mount Doom with the One Ring.",
    "The Green Mile": "The lives of guards on Death Row are affected by one of their charges: a black man accused of child murder and rape, yet who has a mysterious gift.",
    "The Godfather Part II": "The early life and career of Vito Corleone in 1920s New York City is portrayed, while his son, Michael, expands and tightens his grip on the family crime syndicate.",
    "The Silence of the Lambs": "A young F.B.I. cadet must receive the help of an incarcerated and manipulative cannibal killer to help catch another serial killer, a madman who skins his victims.",
    "Goodfellas": "The story of Henry Hill and his life in the mob, covering his relationship with his wife Karen Hill and his mob partners Jimmy Conway and Tommy DeVito in the Italian-American crime syndicate.",
    "Saving Private Ryan": "Following the Normandy Landings, a group of U.S. soldiers go behind enemy lines to retrieve a paratrooper whose brothers have been killed in action.",
    "The Departed": "An undercover cop and a mole in the police attempt to identify each other while infiltrating an Irish gang in South Boston.",
    "Gladiator": "A former Roman General sets out to exact vengeance against the corrupt emperor who murdered his family and sent him into slavery.",
    "Titanic": "A seventeen-year-old aristocrat falls in love with a kind but poor artist aboard the luxurious, ill-fated R.M.S. Titanic."

}



class MovieRecommender:
    def __init__(self, movies, movie_descriptions):
        self.movies = movies
        self.movie_descriptions = movie_descriptions
        self.vectorizer = TfidfVectorizer(stop_words='english')
        self.movie_matrix = self._create_movie_matrix()

    def _create_movie_matrix(self):
        movie_texts = [self.movie_descriptions[movie["title"]] for movie in self.movies]
        movie_matrix = self.vectorizer.fit_transform(movie_texts)
        return movie_matrix

    def recommend_movies(self, user_input, top_n=3):
        user_input_vector = self.vectorizer.transform([user_input])
        print(user_input_vector)
        similarities = cosine_similarity(user_input_vector, self.movie_matrix)
        print(similarities)

        # similarities.argsort()：这一部分是用来获取相似度数组
        # similarities
        # 中元素的索引按照值从小到大排序后的结果。.argsort()
        # 返回的是按照元素大小排序后的索引数组，即原数组中最小元素的索引在第一个位置，依次类推。
        #
        # [0]：由于
        # similarities
        # 是一个二维数组，我们需要获取其第一行（索引为0），因为我们只有一个用户输入，所以只有一行。
        #
        # [-top_n:]：这一部分是用来获取排序后的最后
        # top_n
        # 个元素的索引。在数组切片操作中，[-top_n:] 表示从倒数第
        # top_n
        # 个元素开始直到最后一个元素。
        #
        # [::-1]：最后一步是将获取到的索引数组进行逆序操作，即将其倒置。这是因为在
        # argsort()
        # 返回的索引数组中，索引越靠后，相似度越大。
        similar_movies_indices = similarities.argsort()[0][-top_n:][::-1]
        recommended_movies = [self.movies[i]["title"] for i in similar_movies_indices]
        return recommended_movies


# 初始化推荐系统
recommender = MovieRecommender(movies, movie_descriptions)

# 用户输入
# user_input = "action-packed movie with intense fight scenes"
user_input = "gritty crime drama with a suspenseful cat-and-mouse game between the detective and the criminal mastermind"




# 获取推荐电影
recommended_movies = recommender.recommend_movies(user_input)
print("Recommended Movies:")
for movie in recommended_movies:
    print("-", movie)
